import argparse

import torch
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import utils
from Final.models import generator

parser = argparse.ArgumentParser()
parser.add_argument("--batch_size", type=int, default=128, help="size of the batches")
parser.add_argument("--noise_dim", type=int, default=100, help="dimensionality of the noise")
parser.add_argument("--generator_path", type=str, default="./models/generator.pth", help="name of generator")
parser.add_argument("--normal", type=int, default=0, help="whether to normal the noise")
opt = parser.parse_args()
print(opt)


def Validate(dataloader):
    G = generator(opt.noise_dim, 1 * 56 * 56)
    G.load_state_dict(torch.load(opt.generator_path))
    G.cpu()
    images, labels = utils.make_train_set(model=G, batch_size=opt.batch_size, noise_dim=opt.noise_dim,
                                          dataloader=dataloader, normal=opt.normal)
    correct = 0
    for i in range(400):
        train_images, test_images, train_labels, test_labels = train_test_split(images, labels, train_size=399,
                                                                                random_state=None)
        model = KNeighborsClassifier()
        model.fit(train_images, train_labels)
        predict = model.predict(test_images)
        if predict == test_labels:
            correct += 1
    acc = correct
    print(acc)


dataloader = utils.loadData(batch_size=opt.batch_size)
Validate(dataloader)
